#include "GBufferProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "Render/Texture/CustomTexture.h"
#include "LocalShader.h"
#include "Render/VulkanTool.h"
#include "RayTracingMath.h"
#include "LocalScene.h"
#include "ResourceProcess.h"

void GBufferProcess::OnInit()
{
	mLocalScene = static_cast<LocalScene*>(mGameScene);
	mResourceProcess = mLocalScene->GetProcess<ResourceProcess>("ResourceProcess");

	PerFrameUbo perFrameUbo = mResourceProcess->mPerFrameUbo;

	ENQUEUE_RENDER_COMMAND(BuildModel)([this, perFrameUbo](VkCommandBuffer commandBuffer)
		{
			RefCountPtr<RHIBuffer> perFrameUbo = GetBuffer("ResourceProcess", "perFrameUbo");
			RefCountPtr<RHIBuffer> materialBuffer = GetBuffer("ResourceProcess", "materials");

			mDescriptor = RHIDescriptor::Builder()
				.PushShader<GBufferVertexShader>()
				.PushShader<GBufferFragmentShader>()
				.SetDescriptorCount(0, 1)
				.Build();

			mDescriptor->SetValue("perFrameUBO", perFrameUbo);
			mDescriptor->SetValue("materials", materialBuffer);
			mDescriptorMask = mDescriptor->AllocateDescriptor();

			ColorRenderTarget::Builder gBuffer1Builder;
			gBuffer1Builder
				.SetFormat(VK_FORMAT_R8G8B8A8_UNORM)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
			mGBuffer1[0] = gBuffer1Builder.Build(this, "gBuffer1_s");
			mGBuffer1[1] = gBuffer1Builder.Build(this, "gBuffer1_d");

			ColorRenderTarget::Builder gBuffer2Builder;
			gBuffer2Builder
				.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
			mGBuffer2[0] = gBuffer2Builder.Build(this, "gBuffer2_s");
			mGBuffer2[1] = gBuffer2Builder.Build(this, "gBuffer2_d");

			ColorRenderTarget::Builder gBuffer3Builder;
			gBuffer3Builder
				.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.SetClearColorValue({ 0.f, 0.f, 0.f, -1.f });
			mGBuffer3[0] = gBuffer3Builder.Build(this, "gBuffer3_s");
			mGBuffer3[1] = gBuffer3Builder.Build(this, "gBuffer3_d");

			DepthRenderTarget::Builder gBufferDepthBuilder;
			gBufferDepthBuilder
				.SetUsage(VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetAspectMask(VK_IMAGE_ASPECT_DEPTH_BIT)
				.SetFormat(VK_FORMAT_D32_SFLOAT_S8_UINT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
			mGBufferDepth[0] = gBufferDepthBuilder.Build(this, "gBufferDepth_s");
			mGBufferDepth[1] = gBufferDepthBuilder.Build(this, "gBufferDepth_d");

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mGBuffer1[0])
				.AddColorRenderTarget(mGBuffer2[0])
				.AddColorRenderTarget(mGBuffer3[0])
				.SetDepthRenderTarget(mGBufferDepth[0])
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetCullMode(VK_CULL_MODE_NONE)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mGBuffer1[0])
				.AddRenderTarget(mGBuffer2[0])
				.AddRenderTarget(mGBuffer3[0])
				.AddRenderTarget(mGBufferDepth[0])
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build();

			mFrameBuffer[0] = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();

			mFrameBuffer[1] = RHIFrameBuffer::Builder()
				.AddRenderTarget(mGBuffer1[1])
				.AddRenderTarget(mGBuffer2[1])
				.AddRenderTarget(mGBuffer3[1])
				.AddRenderTarget(mGBufferDepth[1])
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void GBufferProcess::OnTick(float deltaTime)
{
	std::vector<GBufferPushConstants> gBufferPushConstants = mResourceProcess->mGBufferPushConstants;

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, gBufferPushConstants](VkCommandBuffer commandBuffer)
		{

			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer[mResourceProcess->mCurrMask]->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mDescriptorMask.GetDescriptorSet(), 0, nullptr);

			for (uint32 i = 0; i < mResourceProcess->mOpaqueEntitys.size(); ++i)
			{
				LocalEntity* opaqueEntity = mResourceProcess->mOpaqueEntitys[i];

				VkBuffer vertexBuffer = opaqueEntity->mVertexBuffer->GetBuffer();
				VkDeviceSize offsets = 0;
				vkCmdBindVertexBuffers(commandBuffer, 0, 1, &vertexBuffer, &offsets);

				VkBuffer indexBuffer = opaqueEntity->mIndexBuffer->GetBuffer();
				vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);

				vkCmdPushConstants(commandBuffer, mPipeline->GetPipelineLayout(), VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT, 0, sizeof(GBufferPushConstants), &gBufferPushConstants[i]);

				vkCmdDrawIndexed(commandBuffer, opaqueEntity->mIndexBuffer->GetIndexCount(), 1, 0, 0, 0);
			}

			vkCmdEndRenderPass(commandBuffer);
		});
}

void GBufferProcess::OnResize()
{
	ColorRenderTarget::Builder gBuffer1Builder;
	gBuffer1Builder
		.SetFormat(VK_FORMAT_R8G8B8A8_UNORM)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
	mGBuffer1[0] = gBuffer1Builder.Build(this, "gBuffer1_s");
	mGBuffer1[1] = gBuffer1Builder.Build(this, "gBuffer1_d");

	ColorRenderTarget::Builder gBuffer2Builder;
	gBuffer2Builder
		.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
	mGBuffer2[0] = gBuffer2Builder.Build(this, "gBuffer2_s");
	mGBuffer2[1] = gBuffer2Builder.Build(this, "gBuffer2_d");

	ColorRenderTarget::Builder gBuffer3Builder;
	gBuffer3Builder
		.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.SetClearColorValue({ 0.f, 0.f, 0.f, -1.f });
	mGBuffer3[0] = gBuffer3Builder.Build(this, "gBuffer3_s");
	mGBuffer3[1] = gBuffer3Builder.Build(this, "gBuffer3_d");

	DepthRenderTarget::Builder gBufferDepthBuilder;
	gBufferDepthBuilder
		.SetUsage(VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetAspectMask(VK_IMAGE_ASPECT_DEPTH_BIT)
		.SetFormat(VK_FORMAT_D32_SFLOAT_S8_UINT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
	mGBufferDepth[0] = gBufferDepthBuilder.Build(this, "gBufferDepth_s");
	mGBufferDepth[1] = gBufferDepthBuilder.Build(this, "gBufferDepth_d");

	mFrameBuffer[0] = RHIFrameBuffer::Builder()
		.AddRenderTarget(mGBuffer1[0])
		.AddRenderTarget(mGBuffer2[0])
		.AddRenderTarget(mGBuffer3[0])
		.AddRenderTarget(mGBufferDepth[0])
		.SetRenderPass(mRenderPass)
		.Build();

	mFrameBuffer[1] = RHIFrameBuffer::Builder()
		.AddRenderTarget(mGBuffer1[1])
		.AddRenderTarget(mGBuffer2[1])
		.AddRenderTarget(mGBuffer3[1])
		.AddRenderTarget(mGBufferDepth[1])
		.SetRenderPass(mRenderPass)
		.Build();
}


